import time

from .hook import BeforeIterHook, AfterIterHook, BeforeEpochHook, AfterEpochHook


class IterTimerHook(BeforeIterHook, AfterIterHook):

    def __init__(self, name: str = "IterTimerHook"):
        super().__init__(name)

    def before_iter(self, container):
        self.start = time.time()

    def after_iter(self, container):
        duration = int((time.time() - self.start) * 1000)
        self.log("%dms/iter" % duration, container)


class EpochTimerHook(BeforeEpochHook, AfterEpochHook):

    def __init__(self, name: str = "EpochTimerHook"):
        super().__init__(name)

    def before_epoch(self, container):
        self.start = time.time()

    def after_epoch(self, container):
        duration = int(time.time() - self.start)
        self.log("%ds/epoch" % duration, container)
